#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>
#include <libspe2.h>
#include <math.h>
#include <pthread.h>

#include "fish.h"
#include "flocking.h"

/* Flocking Parameters */
const float C0 = 0.00000001; // weight of vector away from neighbors within crowding radius
//const float C1 = 2000.0; // weight of vector towards center of world
const float C1 = 0.0;
const float C2 = 1.0; // weight of velocity correcting vectors towards avg. flocking velocity
const float C3 = 0.0; // weight of vector towards center of gravity
const float C4 = 0.05; // weight of random vector
const float C5 = 0.1; // weight of vector pointing to cog of local flockmates
const float R_CROWD2 = 0.00390625; // square of crowding radius
const float R_SIGHT2 = 0.015625; // square of sighting radius
const float R_SIGHT = 0.125; // sighting radius
const float BIRD_MASS = 1.0; // bird mass
const float R = 1.0; // repulsion constant
const float M = 1.0; // constant that converts vectors to acceleration
const float START_CONST = 0.5;
const float SIGHT_ANGLE = 200.0 * M_PI / 180.0;
const float SCALE_FACTOR = 1.0;
float v0x, v0y, v1x, v1y, v2x, v2y, v3x, v3y, v4x, v4y, v5x, v5y, vx, vy, ax, ay; // vectors
float cogx, cogy; // center of gravity coordinates
float local_cogx, local_cogy;
float tx, ty; // translation vectors for the moving grid

float max_norm;

FILE *output_fp; /* Pointer to output file */
FILE *input_fp; /* Pointer to input file */
int outputp;
int inputp;

int n_fish = 10; /* Other defaults in fish-lib.c. */
//fish_t *fish; /* Global array of fish */

fish_t* sorted;
fish_t *fish[2][BUCKETS_Y][BUCKETS_X];    /* Buckets of fish */
int fish_bucket_size[BUCKETS_Y][BUCKETS_X]; /* Capacities of buckets */
int fish_bucket_n[2][BUCKETS_Y][BUCKETS_X] __attribute__((aligned(16))); /* How many fish in each bucket */

/* SPE PROGRAM STUFF */
spe_params_t spe_params __attribute__((aligned(16)));
spe_auxinfo_t spe_auxinfo[BUCKETS_Y][BUCKETS_X] __attribute__((aligned(16)));
spe_program_handle_t *prog;
pthread_t threads[BUCKETS_Y][BUCKETS_X];
spe_gang_context_ptr_t hells_angels;
void* run_sim_spe(void*);


typedef struct {
  spe_context_ptr_t spe;
  int bucket_number;
} thread_args_t;
thread_args_t thread_args[BUCKETS_Y][BUCKETS_X];  

struct Timer total_timer;

void init_fish();
void interact_fish(int t);
void interact_fish_bucket(int t, int n, int m);
void bounce_fish (fish_t* fish);
void move_fish (float dt, int t);
float compute_norm (int t);
void output_fish(FILE* output_fp, const float curr_time,
    const float dt, const int t);
int fish_bucket_x(fish_t const* f);
int fish_bucket_y(fish_t const* f);
void sort_fish(int);
void init_spes();
void sort_fish_by_bucket(int, int, int);

float dt;

int
main(int argc, char **argv)
{
  //printf("Starting Program\n");
  float curr_time;
  float output_time;
  int t, steps, i, j, ret;
  spe_mbox_msg_t spe_action_msg; // stop or go
  spe_mbox_msg_t spe_dt_msg; // for sending them dt
  spe_mbox_msg_t spe_status_msg; // get status from spe
  unsigned int msg_count;

  get_options(argc, argv);
  //srand48(clock());

  if (output_filename) {
    outputp = 1;
    output_fp = fopen(output_filename, "w");
    if (output_fp == NULL) {
      printf("Could not open %s for output\n", output_filename);
      exit(1);
    }
    fprintf(output_fp, "n_fish: %d\n", n_fish);
  }

  if (input_filename) {
    inputp = 1;
    input_fp = fopen(input_filename, "r");
    if (input_fp == NULL) {
      printf("Could not open %s for input\n", input_filename);
      exit(1);
    }
  }
  
  //printf("Initializing Fish...\n");
  init_fish();
  //printf("Initializing SPEs...\n");
  init_spes();

  if (inputp) { 
    fclose(input_fp);
  }

  initialize_timer (&total_timer);
  start_timer(&total_timer);

  for (output_time = 0.0, curr_time = 0.0, steps = 0, t = 0;
       curr_time <= end_time && steps < max_steps;
       curr_time += dt, ++steps, t ^= 1) {

    // Basically our plan now is to send 1 of two messages to the spes:
    // - Go message (1) - means get dt from your mailbox, and go
    // - Stop messages (0) - means exit
    // Upon receival of a Go message, the spe will grab one more message
    // from it's mailbox, which should be dt for that round.
    // When the SPEs are done, they will send a done message back
    //    printf("Sending GO_MSG to all SPEs, t = %d\n", t);
    spe_action_msg.i = GO_MSG;
    spe_dt_msg.f = dt;
    
    // send GO message and dt
    for( i = 0; i<BUCKETS_Y; i++ ) {
      for( j = 0; j<BUCKETS_X; j++ ) {
	if( -1 == spe_in_mbox_write( thread_args[i][j].spe, &(spe_action_msg.i), 1, SPE_MBOX_ANY_NONBLOCKING ) ) {
	  perror("spe_in_mbox_write()");
	  exit(1);
	}
	if( -1 == spe_in_mbox_write( thread_args[i][j].spe, &(spe_dt_msg.i), 1, SPE_MBOX_ANY_NONBLOCKING ) ) {
	  perror("spe_in_mbox_write()");
	  exit(1);
	}
      }
    }

    /* We only output once every output_interval time unit, at most.
     * Without that restriction, we can easily create a huge output
     * file.  Printing a record for ten fish takes about 300 bytes,
     * so for every 1000 steps, we could dump 300K of info.  Now
     * scale the number of fish by 1000...
     */
    if (outputp && curr_time >= output_time) {
      output_fish(output_fp, curr_time, dt, t);
      output_time = curr_time + output_interval;
    }
    /*
    printf("fish_bucket_n[0] array:\n\t\t%d\t%d\t%d\n\t\t%d\t%d\t%d\nfish_bucket_n[1] array:\n\t\t%d\t%d\t%d\n\t\t%d\t%d\t%d\n", 
	   fish_bucket_n[0][0][0],
	   fish_bucket_n[0][0][1],
	   fish_bucket_n[0][0][2],
	   fish_bucket_n[0][1][0],
	   fish_bucket_n[0][1][1],
	   fish_bucket_n[0][1][2],
	   fish_bucket_n[1][0][0],
	   fish_bucket_n[1][0][1],
	   fish_bucket_n[1][0][2],
	   fish_bucket_n[1][1][0],
	   fish_bucket_n[1][1][1],
	   fish_bucket_n[1][1][2] );
    */


    // wait for SPEs to send back done message
    msg_count = 0;
    max_norm = 0.0; // cuz sort_fish_by_bucket needs to use this global
    cogx = cogy = 0.0; // sort_fish_by_bucket also tracks this                                                                                                   
    memset( fish_bucket_n[t^1], 0, BUCKETS_X*BUCKETS_Y*sizeof(int) );

    while( msg_count != BUCKETS_X*BUCKETS_Y ) {
      for( i = 0; i<BUCKETS_Y; i++ ) {
	for( j = 0; j<BUCKETS_X; j++ ) {
	  if( spe_out_mbox_status( thread_args[i][j].spe ) ) {
	    spe_out_mbox_read( thread_args[i][j].spe, &(spe_status_msg.i), 1);
	    if( spe_status_msg.i == DONE_MSG ) {
	      msg_count++;
	      sort_fish_by_bucket(t, i, j); // don't wait, immediately sort the fish we have
	    }
	    else {
	      printf("Error: Expecting DONE_MSG, received %d\n", spe_status_msg.i);
	      exit(1);
	    }
	  }
	}
      }
    }

    //printf("All SPEs are done\n");
    // IDEA: sort buckets as they arrive instead of once they all arrive    
    //sort_fish(t);

    // compute dt                                                                                                                                            
    dt = max_norm_change / max_norm;
    dt = f_max(dt, min_dt);
    dt = f_min(dt, max_dt);    

    cogx /= n_fish;
    cogy /= n_fish;
    
    tx = scale_coord(cogx) - 0.5;
    ty = scale_coord(cogy) - 0.5;
  }

  // Is this really the right time to stop the timer? Technically we have the data
  stop_timer(&total_timer);

  // Send STOP message to all the spes.
  spe_action_msg.i = STOP_MSG;
    
  // send STOP message
  for( i = 0; i<BUCKETS_Y; i++ ) {
    for( j = 0; j<BUCKETS_X; j++ ) {
      if( -1 == spe_in_mbox_write( thread_args[i][j].spe, &(spe_action_msg.i), 1, SPE_MBOX_ANY_NONBLOCKING ) ) {
	perror("spe_in_mbox_write()");
	exit(1);
      }
    }
  }

  // Join on threads to wait for them to stop
  for(i = 0; i < BUCKETS_Y; i++) {
    for(j = 0; j < BUCKETS_X; j++) {
      //printf("Joining on (%d, %d)\n", i, j);
      pthread_join( threads[i][j], NULL );
      //printf("Thread (%d, %d) finished\n", i, j);
      ret = spe_context_destroy(thread_args[i][j].spe);
      if( ret ) {
	perror("spe_context_destroy");
	exit(1);
      }
    }
  }

  spe_gang_context_destroy(hells_angels);  

  if (outputp) {
    output_fish(output_fp, curr_time, dt, t);
    fclose(output_fp);
  }


  // first number is the simulation version
  printf("%d\t%d\t%g\t%g\t%d\t%d\t%g\n", 7, n_fish,
         curr_time, end_time, steps, max_steps, timer_duration(total_timer));
  /*
  printf("Ended at %g (%g), %d (%d) steps\n",
      curr_time, end_time, steps, max_steps);
  printf("Total time taken: %g\n", timer_duration(total_timer));
  */
  return 0;
}

void
sort_fish_by_bucket(int t, int i, int j)
{
  int x, y, k, nx, ny, tp1;
  
  tp1 = t^1;
  
  for(k = 0; k < fish_bucket_n[t][i][j]; k++) {
    fish_t *__restrict__ f = &fish[t][i][j][k];
    
    cogx += f->x;
    cogy += f->y;
    
    max_norm = f_max(max_norm, fabs(fish[t][i][j][k].vx));
    max_norm = f_max(max_norm, fabs(fish[t][i][j][k].vy));
    max_norm = f_max(max_norm, fabs(fish[t][i][j][k].ax));
    max_norm = f_max(max_norm, fabs(fish[t][i][j][k].ay));

    nx = fish_bucket_x(f);
    ny = fish_bucket_y(f);

    fish[tp1][ny][nx][fish_bucket_n[tp1][ny][nx]++]
      = fish[t][i][j][k]; 

    // Just as a test, let's memset the previous fish to 0 - test passed
    //memset( fish[t][i][j], 0, fish_bucket_n[t][i][j]*sizeof(fish_t) );
  }
}


void
sort_fish(int t)
{
  int x, y, i, j, k, nx, ny, tp1;
  float max_norm;
  
  tp1 = t^1;

  cogx = cogy = 0.0;

  //printf("memsetting\n");
  memset( fish_bucket_n[tp1], 0, BUCKETS_X*BUCKETS_Y*sizeof(int) );

  max_norm = 0.0;

  for(i = 0; i < BUCKETS_Y; i++){
    for(j = 0; j < BUCKETS_X; j++){
      for(k = 0; k < fish_bucket_n[t][i][j]; k++)
      {
	fish_t *__restrict__ f = &fish[t][i][j][k];

	cogx += f->x;
	cogy += f->y;

	max_norm = f_max(max_norm, fabs(fish[t][i][j][k].vx));
	max_norm = f_max(max_norm, fabs(fish[t][i][j][k].vy));
	max_norm = f_max(max_norm, fabs(fish[t][i][j][k].ax));
	max_norm = f_max(max_norm, fabs(fish[t][i][j][k].ay));

        nx = fish_bucket_x(f);
        ny = fish_bucket_y(f);

        fish[tp1][ny][nx][fish_bucket_n[tp1][ny][nx]++]
          = fish[t][i][j][k]; 
      }
      // Just as a test, let's memset the previous fish to 0 - test passed
      //memset( fish[t][i][j], 0, fish_bucket_n[t][i][j]*sizeof(fish_t) );
    }
  }

  // compute dt                                                                                                                                            
  dt = max_norm_change / max_norm;
  dt = f_max(dt, min_dt);
  dt = f_min(dt, max_dt);

  cogx /= n_fish;
  cogy /= n_fish;
 
  tx = scale_coord(cogx) - 0.5;
  ty = scale_coord(cogy) - 0.5;
}


void
init_spes()
{
  int i, j, ret;
  hells_angels = spe_gang_context_create(0);

  if( NULL == hells_angels ) {
    perror("spe_gang_context_create");
    exit(1);
  }

  prog = spe_image_open("./fish_spe.elf");
  if( !prog ) {
    perror("spe_image_open");
    exit(1);
  }

  spe_params.fish_bucket_n_ea = (unsigned long)fish_bucket_n;

  /* initialize variables */
  for( i = 0; i < BUCKETS_Y; i++ ) {
    for( j = 0; j < BUCKETS_X; j++ ) {
      thread_args[i][j].spe = spe_context_create(0, hells_angels);
      if( !thread_args[i][j].spe ) {
	perror("spe_context_create");
	exit(1);
      }

      thread_args[i][j].bucket_number = i*BUCKETS_X + j;
      
      ret = spe_program_load( thread_args[i][j].spe, prog );
      if( ret ) {
	perror("spe_program_load");
	exit(1);
      }

      spe_auxinfo[i][j].bucket_number = thread_args[i][j].bucket_number;
      spe_auxinfo[i][j].random_seed = drand48();
      spe_params.fish_ea[0][i][j] = (unsigned long)(fish[0][i][j]);
      spe_params.fish_ea[1][i][j] = (unsigned long)(fish[1][i][j]);
    }
  }
 
  //printf("init_spes(): Creating Pthreads\n");

  for(i = 0; i < BUCKETS_Y; i++){
    for(j = 0; j < BUCKETS_X; j++){
      //printf("init_spes(): Creating thread:\n\tcoordinate: (%d, %d)\n\tbucket number: %d\n", i, j, thread_args[i][j].bucket_number); 
      ret = pthread_create( &threads[i][j], NULL, run_sim_spe, &thread_args[i][j] );
      if( ret ) {
	perror("pthread_create");
	exit(1);
      }
    }
  }

}


/* Thread function to run our program on an spe
 */
void*
run_sim_spe(void* thread_arg)
{
  int ret;
  thread_args_t* args = (thread_args_t*)thread_arg;
  unsigned int entry;
  spe_stop_info_t stop_info;

  int n = (args->bucket_number)/BUCKETS_X;
  int m = (args->bucket_number)%BUCKETS_X;

  //printf("run_sim_spe(): Booting Context (%d, %d). Bucket # %d\n", n, m, args->bucket_number);

  entry = SPE_DEFAULT_ENTRY;

  //printf("run_sim_spe(): Sending (%d, %d):\n\tspe_auxinfo[%d][%d].bucket_number = %d\n\t&spe_params: %lu\n\t&spe_auxinfo[%d][%d]: %lu\n\tfish_bucket_n: %lu\n", n, m, n, m, spe_auxinfo[n][m].bucket_number, &spe_params, n, m, &(spe_auxinfo[n][m]), fish_bucket_n );
  ret = spe_context_run( args->spe, &entry, 0, &spe_params, &(spe_auxinfo[n][m]), &stop_info );

  if( ret < 0 ) {
    perror("spe_context_run");
    return NULL;
  }

  return NULL;
}


/* Allocate and initialize the fish positions / velocities / accelerations. */
void
init_fish()
{
  int i, j, nx, ny;
  float xf, yf;
  fish_t *fishes;
  float max_norm;

  /* I'll just do this the easy way - allocate and initialize them all at
    once, and then split them into buckets. This way is easier to code, and
    results in fewer calls into stdlib malloc/realloc etc. */

  if(!(fishes = (fish_t*)malloc(n_fish*sizeof(fish_t)))){
    perror("malloc");
    exit(-1);
  }
  
  if(uniformp){
    /*srand48(clock());*/
    srand48(0);
  }
  
  cogx = cogy = 0;

  for(i = 0; i < n_fish; i++){
    if(inputp){
      fscanf(input_fp, "  %d: (%f, %f)", &j, &xf, &yf);
      fishes[i].x = unscale_coord(xf);
      fishes[i].y = unscale_coord(yf);
    }else if(uniformp){
      fishes[i].x = unscale_coord(drand48());
      fishes[i].y = unscale_coord(drand48());
    }else{
      float angle = i * (2.0 * M_PI / n_fish);
      fishes[i].x = unscale_coord(0.5 * cos(angle) + 0.5);
      fishes[i].y = unscale_coord(0.5 * sin(angle) + 0.5);
    }
    cogx += fishes[i].x;
    cogy += fishes[i].y;
  }

  cogx /= n_fish;
  cogy /= n_fish;
  tx = scale_coord(cogx) - 0.5;
  ty = scale_coord(cogy) - 0.5;

  for(i = 0; i < n_fish; i++){
    fishes[i].ax = fishes[i].ay = 0.0;
    if(inputp){
      fscanf(input_fp, "  %d: (%f, %f)", &j, &xf, &yf);
      fishes[i].vx = xf;
      fishes[i].vy = yf;
    }else{
      fishes[i].vx = START_CONST*(drand48() - 0.5);
      fishes[i].vy = START_CONST*(drand48() - 0.5);
    }
  }

  for(i = 0; i < BUCKETS_Y; i++){ // rows (y-direction)
    for(j = 0; j < BUCKETS_X; j++) // columns (x-direction)
    {
      fish_bucket_size[i][j] = 1;
      fish_bucket_n[0][i][j] = 0;
      fish_bucket_n[1][i][j] = 0;
    }
  }

  for(i = 0; i < n_fish; i++){
    nx = fish_bucket_x(&fishes[i]);
    ny = fish_bucket_y(&fishes[i]);
    fish_bucket_size[ny][nx]++; 
  }

  for(i = 0; i < BUCKETS_Y; i++){
    for(j = 0; j < BUCKETS_X; j++) {
      // int size = sizeof(fish_t)*fish_bucket_size[i][j];
      // use memalign to allocate an aligned piece of
      // memory... to 16 bytes for DMA transfers
      int size = sizeof(fish_t)*n_fish;
      if(!(fish[0][i][j] = (fish_t*)memalign(16, size)) ||
	 !(fish[1][i][j] = (fish_t*)memalign(16, size)))
      {
        perror("memalign");
        exit(-1);
      }
    }
  }

  max_norm = 0.0;

  // sort the fish into their corect buckets
  for(i = 0; i < n_fish; i++){
    nx = fish_bucket_x(&fishes[i]);
    ny = fish_bucket_y(&fishes[i]);
    fish[0][ny][nx][ fish_bucket_n[0][ny][nx]++ ] = fishes[i];

    // compute max_norm for first iteration
    max_norm = f_max(max_norm, fabs(fishes[i].vx));
    max_norm = f_max(max_norm, fabs(fishes[i].vy));
  }

  // compute dt
  dt = max_norm_change / max_norm;
  dt = f_max(dt, min_dt);
  dt = f_min(dt, max_dt);

  free(fishes);

  sorted = malloc(n_fish*sizeof(fish_t));
  if(!sorted){
    perror("malloc");
    exit(-1);
  }
}

/*
  Dump out all the fishies (and their center of gravity)
  in a format that the viewer understands.
*/


void
output_fish(FILE* output_fp, const float curr_time,
    const float dt, const int t)
{
  int i, j, k, l, x, y;
  
  l = 0;
  for(i = 0; i < BUCKETS_Y; i++){
    for(j = 0; j < BUCKETS_X; j++){
      for(k = 0; k < fish_bucket_n[t][i][j]; k++){
        sorted[l] = fish[t][i][j][k];
	l++;
      }
    }
  }

  fprintf(output_fp, "%.5g (%.5g):\n", curr_time, dt);
  for (i = 0; i < n_fish; ++i) {
    fprintf(output_fp, "  %d: (%g, %g)\n", i,
        scale_coord(sorted[i].x), scale_coord(sorted[i].y));
  }
  if (outvelp) {
    for (i =0; i< n_fish; ++i) {
      fprintf(output_fp, "  %d: (%g, %g)\n", i, sorted[i].vx, sorted[i].vy);
    }
  }
  fprintf(output_fp, "  cg: (%g, %g)\n", 1.5, 1.5);
}

int
fish_bucket_x(fish_t const* f)
{
  //return floor(fmod(scale_coord(f->x)-tx+1,1)*BUCKETS); 
  return floor(fmod(scale_coord(f->x)-tx+1,1)*BUCKETS_X); 
}

int
fish_bucket_y(fish_t const* f)
{
  //return floor(fmod(scale_coord(f->y)-ty+1,1)*BUCKETS); 
  return floor(fmod(scale_coord(f->y)-ty+1,1)*BUCKETS_Y); 
}


